#include <iostream>
#include <limits.h>

#include "atom.h"
#include "bytecode.h"
#include "generalizer.h"
#include "globals.h"
#include "predicates.h"
#include "prover.h"
#include "reference.h"
#include "rules.h"
#include "serialize.h"
#include "tree.h"
#include "tuple.h"

// Re-entrancy flag
__thread int recursionLevel = 0;

// Caches prover factory for system or bytecode-compiled predicates
void optimizeProverBytecode (Bytecode &bytecode, Node &call, System &system) {
  if(typeid(call) != typeid(Reference)) {
    Prototype proto = Prototype::get (call);

    if(proto != Prototype ()) {
      ProverFactoryBase *factory = Prover::findSystemPredicate (proto);
      if(factory) {
        bytecode.o2.systemProverFactory = factory;
        bytecode.instruction = PROVERSYS;
        return;
      }

      Snippet *snippet = &system.rules.findSnippet (proto);
      if(snippet) {
        bytecode.o2.snippet = snippet;
        bytecode.instruction = PROVERBCD;
        return;
      }

      bytecode.instruction = PROVERPRO; // Can be binded to snippet later
    }
  }
}

// Returns the last variable assigned in code
void vectorizeBytecode (Node &node, Snippet &snippet, System &system) {
  Tuple &codeTuple = static_cast<Tuple&> (*node.final ());

  Bytecode bytecode;
  std::string tag = codeTuple.nodes[1].display ();

  // Parse the operands automatically if they are numbers
  Node *o1 = codeTuple.nodes[2].final ();
  Node *o2 = codeTuple.nodes[3].final ();

  if(typeid(*o1) == typeid(Number)) {
    bytecode.o1.var = static_cast<Number&> (*o1).value;
    if(bytecode.o1.var >= snippet.nVariables)
      snippet.nVariables = bytecode.o1.var + 1;
  }

  if(typeid(*o2) == typeid(Number))
    bytecode.o2.n = static_cast<Number&> (*o2).value;

  if(tag == "BIND")
    bytecode.instruction = BIND_____;
  else if(tag == "BIND.UNDO")
    bytecode.instruction = BINDUNDO_;
  else if(tag == "CALL")
    bytecode.instruction = CALL_____;
  else if(tag == "CUT")
    bytecode.instruction = CUT______;
  else if(tag == "CUT.FAIL")
    bytecode.instruction = CUTFAIL__;
  else if(tag == "CUT.POINT") {
    bytecode.instruction = CUTPOINT_;
    if(bytecode.o1.var >= snippet.nCutPoints)
      snippet.nCutPoints = bytecode.o1.var + 1;
  } else if(tag == "EXEC")
    bytecode.instruction = EXEC_____;
  else if(tag == "FAIL")
    bytecode.instruction = FAIL_____;
  else if(tag == "JUMP")
    bytecode.instruction = JUMP_____;
  else if(tag == "JUMP.FAIL")
    bytecode.instruction = JUMPFAIL_;
  else if(tag == "NEW.REF")
    bytecode.instruction = NEWREF___;
  else if(tag == "PROVEN")
    bytecode.instruction = PROVEN___;
  else if(tag == "PROVER.POP")
    bytecode.instruction = PROVERPOP;
  else if(tag == "PROVER.PUSH") {
    bytecode.instruction = PROVERPSH;
    bytecode.o2.constant = o2;
    optimizeProverBytecode (bytecode, *o2, system);
  } else if(tag == "PROVER.TRY")
    bytecode.instruction = PROVERTRY;
  else if(tag == "RETURN")
    bytecode.instruction = RETURN___;
  else if(tag == "SET.CONST") {
    bytecode.instruction = SETCONST_;
    snippet.cascadeNodes.insert (bytecode.o2.constant = o2->final ());
  } else if(tag == "SET.TOPROVE")
    bytecode.instruction = SETTOPROV;
  else if(tag == "SET.VAR")
    bytecode.instruction = SETVAR___;
  else if(tag == "TUPLE.APPEND")
    bytecode.instruction = TUPLEAPP_;
  else if(tag == "TUPLE.FROM")
    bytecode.instruction = TUPLEFROM;
  else if(tag == "TUPLE.NEW")
    bytecode.instruction = TUPLENEW_;
  else if(tag == "TUPLE.NTH")
    bytecode.instruction = TUPLENTH_;
  else if(tag == "TREE.LR")
    bytecode.instruction = TREELR___;
  else if(tag == "TREE.OP") {
    bytecode.instruction = TREEOP___;
    std::string s = o2->display ();
    bytecode.o2.op[0] = s[0];
    bytecode.o2.op[1] = (s.length () > 1) ? s[1] : 0;
  } else
    std::cerr << "[FATAL] Unknown bytecode " << tag << std::endl;

  snippet.codes.push_back (bytecode);
}

bool vectorizeBytecodes (Node &result, Snippet &snippet, System &system) {
  std::vector<pNode> array =
      static_cast<Tuple*> (result.final ())->nodes;
  snippet.codes.clear ();
  snippet.codes.reserve (array.size ()); // We have our idealism

  // Convert all elements and add up to the vector
  std::vector<pNode>::iterator iter;
  for(iter = array.begin (); iter < array.end (); iter++)
    vectorizeBytecode (**iter, snippet, system);

  return snippet.nVariables < SHRT_MAX;
}

bool callCompiler (const std::string &predName, Node &rule,
    Snippet &snippet, System &system) {
  bool okay = false;
  Reference *result = new Reference, *input = new Reference;

  Tuple *tuple = new Tuple;
  tuple->nodes.push_back (Atom::create (predName));
  tuple->nodes.push_back (input);
  tuple->nodes.push_back (result);
  pNode node (tuple);

  if(recursionLevel++ < 2) { // limits level of re-entrancy
    Globals globals; // Use a new system to avoid synchronization deadlock
    Trail trail;
    System newSystem (system.rules, globals, trail);

    pNode generalized = Generalizer ().generalize (*node);
    input->bound (rule, trail); // Bind late, avoid it being generalized
    Prover *prover = Prover::create (*generalized, newSystem);
    okay = prover->prove ();
    delete prover;
  }
  recursionLevel--;

  if(!okay)
    return false;

  return vectorizeBytecodes (*result, snippet, system);
}

bool compileExecuteBytecode (
    Prototype &proto, Snippet &snippet, System &system) {
  const std::string compileRoutine = "bytecode.compile";
  pNode node;

  if(proto.type == Tuple::ID) {
    Tuple *tuple = new Tuple;
    tuple->nodes.push_back (proto.functor);
    for(int i = 1; i < proto.u.length; i++)
      tuple ->nodes.push_back (new Reference);
    node = tuple; // Keeps it referenced

    return callCompiler (compileRoutine, *node, snippet, system);
  } else if(proto.type == Atom::ID)
    return callCompiler (compileRoutine, *proto.functor, snippet, system);

  return false;
}

BytecodeExecutor::~BytecodeExecutor () {
  for(std::vector<Prover*>::iterator iter = provers.begin (); iter
      < provers.end (); iter++)
    delete *iter;
}

bool BytecodeExecutor::prove () {
  for(;;) {
    //    std::cerr << (void*) this << " EXECUTING " << ip - snippet.codes.begin ()
    //        << std::endl;

    Bytecode *code = &*ip++;

    switch(code->instruction) {
    case BIND_____: {
      binders.push_back (system.trail.state ());
      result = binders.back ().bind (*locals[code->o1.var],
          *locals[code->o2.n], system.trail);
    }
      break;
    case BINDUNDO_:
      binders.back ().undo (system.trail);
      binders.pop_back ();
      break;
    case CALL_____:
      callees.push_back (ip);
      ip = snippet.codes.begin () + code->o2.n;
      break;
    case CUT______: {
      CutPoint &cutPoint = cutPoints[code->o1.var];
      for(std::vector<Prover*>::iterator iter = provers.begin ()
          + cutPoint.proverIndex; iter < provers.end (); iter++)
        delete *iter;

      binders.resize (cutPoint.binderIndex);
      callees.resize (cutPoint.calleeIndex);
      provers.resize (cutPoint.proverIndex);
    }
      break;
    case CUTFAIL__: {
      CutPoint &cutPoint = cutPoints[code->o1.var];
      ip = snippet.codes.begin () + code->o2.n;
      system.trail.rollback (cutPoint.trailPointer);
    }
      break;
    case CUTPOINT_: {
      CutPoint &cutPoint = cutPoints[code->o1.var];
      cutPoint.trailPointer = system.trail.state ();
      cutPoint.binderIndex = binders.size ();
      cutPoint.calleeIndex = callees.size ();
      cutPoint.proverIndex = provers.size ();
    }
      break;
    case EXEC_____: {
      Prover *prover = Prover::create (*locals[code->o1.var], system);
      result = prover->prove ();
      delete prover;
    }
      break;
    case FAIL_____: // We're all broke
      return false;
    case JUMP_____:
      ip = snippet.codes.begin () + code->o2.n;
      break;
    case JUMPFAIL_:
      if(!result) {
        ip = snippet.codes.begin () + code->o2.n;
        result = true;
      }
      break;
    case NEWREF___:
      locals[code->o1.var] = new Reference;
      break;
    case PROVEN___:
      return true;
    case PROVERBCD: {
      Prover *prover;
      if(code->o2.snippet->valid)
        prover = new BytecodeExecutor (
          *code->o2.snippet, *locals[code->o1.var], system);
      else {
        code->instruction = PROVERPRO; // Finds prover again next time
        prover = Prover::create (*locals[code->o1.var], system);
      }
      provers.push_back (prover);
    }
      break;
    case PROVERPOP:
      delete provers.back ();
      provers.pop_back ();
      break;
    case PROVERPRO: // Doing late better than never
      optimizeProverBytecode (*code, *locals[code->o1.var], system);
    case PROVERPSH:
      provers.push_back (Prover::create (*locals[code->o1.var], system));
      break;
    case PROVERSYS:
      provers.push_back (Prover::createSystemPredicate (
          code->o2.systemProverFactory, *locals[code->o1.var].final (), system)
      );
      break;
    case PROVERTRY:
      result = provers.back ()->prove ();
      break;
    case RETURN___:
      ip = callees.back ();
      callees.pop_back ();
      break;
    case SETCONST_:
      locals[code->o1.var] = code->o2.constant;
      break;
    case SETTOPROV:
      locals[code->o1.var] = &toProve;
      break;
    case SETVAR___:
      locals[code->o1.var] = locals[code->o2.n];
      break;
    case TUPLEAPP_:
      tuple->nodes.push_back (locals[code->o1.var]);
      tryTupleAppend:
      while((*ip).instruction == TUPLEAPP_)
        tuple->nodes.push_back (locals[(*ip++).o1.var]); // This likely repeats
      break;
    case TUPLEFROM: {
      Node *node = locals[code->o1.var].final ();
      if(typeid(*node) == typeid(Tuple))
        tuple = static_cast<Tuple*> (node);
      else
        result = false;
    }
      break;
    case TUPLENEW_:
      locals[code->o1.var] = tuple = new Tuple ();
      tuple->nodes.reserve (code->o2.n);
      goto tryTupleAppend;
    case TUPLENTH_:
      locals[code->o1.var] = tuple->nodes[code->o2.n];
      break;
    case TREELR___:
      left = locals[code->o1.var];
      right = locals[code->o2.n];
      if((*ip).instruction == TREEOP___) { // Most often instruction pair
        code = &*ip++;
    case TREEOP___:
        locals[code->o1.var] = new Tree (code->o2.op, left, right);
      }
    }
  }

  return false;
}

void Snippet::serialize (const Serializer &serializer) {
  ::serialize (serializer, nVariables);
  ::serialize (serializer, toProveVariable);
  ::serialize (serializer, nCutPoints);

  int codeSize = codes.size ();
  ::serialize (serializer, codeSize);
  codes.resize (codeSize);
  for(std::vector<Bytecode>::iterator iter = codes.begin ();
      iter != codes.end (); iter++) {
    ::serialize (serializer, *iter);
    short instruction = (*iter).instruction;

    // Some instructions with pointers are changed
    if(instruction == PROVERBCD || instruction == PROVERSYS)
      instruction = PROVERPSH;
    if(instruction == SETCONST_)
      instruction = 0; // FIXME this is a fail
  }
}

class enableBytecode : public OnceProver {
predicateName("enable.bytecode")

  bool proveOnce () {
    system->rules.bytecodeEnabled = tuple ? isTrue (*tuple->nodes[1]) : true;
    return true;
  }

  Tuple *tuple;
  System *system;
  void set (Atom&, System &s) {
    tuple = 0;
    system = &s;
  }
  void set (Tuple &t, System &s) {
    tuple = &t;
    system = &s;
  }
};

ProverFactoryRegistrar<enableBytecode, Atom::ID> enableBytecodeRegistrar;
ProverFactoryRegistrar<enableBytecode, Tuple::ID, 2>
    enableBytecodeTupleRegistrar;

class isCompiled : public OnceProver {
  predicateName("is.compiled");

  bool proveOnce () {
    Prototype proto = Prototype::get (*tuple->nodes[1]);
    return &system->rules.findSnippet (proto) != 0;
  }

  Tuple *tuple;
  System *system;
  void set (Tuple &t, System &s) {
    tuple = &t;
    system = &s;
  }
};

ProverFactoryRegistrar<isCompiled, Tuple::ID, 2> isCompiledRegistrar;

class setBytecode : public OnceProver {
predicateName("set.bytecode")

  bool proveOnce () {
    bool ret = true;
    Prototype proto = Prototype::get (*tuple->nodes[1]);

    Node &codeTupleNode = *tuple->nodes[2].final ();
    if(typeid(codeTupleNode) == typeid(Tuple)) {
      Snippet snippet;
      ret = vectorizeBytecodes (codeTupleNode, snippet, *system);
      if(ret)
        system->rules.putSnippet (proto, snippet);
    } else
      // Unset bytecode by specifying a nil (or anything other than a tuple)
      system->rules.snippets[proto].invalidate ();

    return ret;
  }

  Tuple *tuple;
  System *system;
  void set (Tuple &t, System &s) {
    tuple = &t;
    system = &s;
  }
};

ProverFactoryRegistrar<setBytecode, Tuple::ID, 3> setBytecodeRegistrar;
